import os
import json
import easyocr
import numpy as np
from PIL import Image
import cv2
import ssl

# 禁用SSL证书验证（注意：这只是为了开发目的，生产环境中应该正确配置证书）
ssl._create_default_https_context = ssl._create_unverified_context

def extract_text_from_image(image_path):
    # 初始化 EasyOCR
    reader = easyocr.Reader(['ch_sim', 'en'])
    
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        return None, None
    
    # 获取图片尺寸
    height, width = img.shape[:2]
    
    # 使用 EasyOCR 识别文字
    results = reader.readtext(img)
    
    # 分类文字（根据字体大小和位置）
    title = []
    describe = []
    
    # 假设标题通常在顶部且字体较大
    for bbox, text, conf in results:
        # 计算文本框的高度
        box_height = abs(bbox[2][1] - bbox[0][1])
        y_position = bbox[0][1]  # 文本的y坐标
        
        # 如果文本框高度较大且在图片上部，认为是标题
        if box_height > height * 0.05 and y_position < height * 0.3:
            title.append(text)
        else:
            describe.append(text)
    
    return ' '.join(title), describe

def process_directory(base_path):
    # 遍历所有子目录
    for dir_name in os.listdir(base_path):
        dir_path = os.path.join(base_path, dir_name)
        if os.path.isdir(dir_path) and not dir_name.startswith('.'):  # 忽略隐藏文件夹
            try:
                # 查找图片文件
                for file_name in os.listdir(dir_path):
                    if file_name.lower().endswith(('.jpg', '.jpeg', '.png')):
                        image_path = os.path.join(dir_path, file_name)
                        
                        # 提取文字
                        title, describe = extract_text_from_image(image_path)
                        
                        if title is not None:
                            # 创建JSON数据
                            data = {
                                "title": title,
                                "describe": describe
                            }
                            
                            # 保存到文本文件
                            output_path = os.path.join(dir_path, 'image.txt')
                            with open(output_path, 'w', encoding='utf-8') as f:
                                json.dump(data, f, ensure_ascii=False, indent=2)
                            
                            print(f'Processed: {dir_name}')
            except Exception as e:
                print(f'Error processing directory {dir_name}: {str(e)}')

if __name__ == "__main__":
    # 处理can目录
    process_directory("can") 